package rbtree

type Node struct {
	key    uint32
	val    any
	left   *Node
	right  *Node
	parent *Node

	color string
}

type RBtree struct {
	root *Node
	nil  *Node
}

func leftRotate(T *RBtree, x *Node) {
	y := x.right
	x.right = y.left
	if y.left != T.nil {
		y.left.parent = x
	}
	y.parent = x.parent
	if x.parent == T.nil {
		T.root = y
	} else if x == x.parent.left {
		x.parent.left = y
	} else {
		x.parent.right = y
	}
	y.left = x
	x.parent = y
}

func rightRotate(T *RBtree, y *Node) {
	x := y.left
	y.left = x.right
	if x.right != T.nil {
		x.right.parent = y
	}
	x.parent = y.parent
	if y.parent == T.nil {
		T.root = x
	} else if y == y.parent.left {
		y.parent.left = x
	} else {
		y.parent.right = x
	}
	x.right = y
	y.parent = x
}

func Insert(T *RBtree, z *Node) int {
	x := T.root
	if x == T.nil {
		T.root = z
		T.root.color = "BLACK"
		return 0
	}
	y := x
	for x != T.nil {
		y = x
		if z.key < x.key {
			x = x.left
		} else if z.key > x.key {
			x = x.right
		} else {
			return -1
		}
	}
	if z.key < y.key {
		y.left = z
	} else {
		y.right = z
	}
	z.parent = y
	z.color = "RED"

	fixup(T, z)
	return 0
}

func Del(T *RBtree, key uint32) int {
	// 找到要删除的节点
	x := T.root
	for x != T.nil && x.key != key {
		if key < x.key {
			x = x.left
		} else {
			x = x.right
		}
	}
	if x == T.nil {
		return -1
	}

	// 情况1: 没有子节点（叶子节点）
	if x.left == T.nil && x.right == T.nil {
		removeNode(T, x)
	} else if (x.left != T.nil && x.right == T.nil) || (x.left == T.nil && x.right != T.nil) {
		y := x
		if x.left != T.nil {
			y = x.left
		} else {
			y = x.right
		}
		removeNode(T, x)
		y.color = "BLACK"
	} else {
		y := findMin(T, x.right)
		x.key = y.key
		x.val = y.val
		// 删除y，y现在指向x原来的右子树中的最小节点
		Del(T, y.key)
	}
	return 0
}

// 移除节点并重新链接树
func removeNode(T *RBtree, z *Node) {
	y := z
	originalColor := y.color
	x := T.nil

	if z.left == T.nil {
		x = z.right
	} else if z.right == T.nil {
		x = z.left
	} else {
		y = findMin(T, z.right)
		originalColor = y.color
		x = y.right
	}

	if x != T.nil {
		x.parent = y.parent
	}

	if y.parent == nil {
		T.root = x
	} else if y == y.parent.left {
		y.parent.left = x
	} else {
		y.parent.right = x
	}

	if y != z {
		z.key = y.key
		z.val = y.val
	}

	// 修复红黑树的性质
	if originalColor == "BLACK" {
		fixup(T, x)
	}
}

// 找到以给定节点为根的子树中的最小节点
func findMin(T *RBtree, x *Node) *Node {
	for x.left != T.nil {
		x = x.left
	}
	return x
}

func Mod(T *RBtree, z *Node) any {
	x := T.root
	for x != T.nil {
		if z.key < x.key {
			x = x.left
		} else if z.key > x.key {
			x = x.right
		} else {
			x.val = z.val
			return z.val
		}
	}
	return nil
}
func fixup(T *RBtree, z *Node) {
	for z != nil && z.color == "RED" {
		if z.parent == nil || z.parent == T.root {
			z.color = "BLACK"
			return
		}
		if z.parent == z.parent.parent.left {
			y := z.parent.parent.right
			if y != nil && y.color == "RED" {
				z.parent.color = "BLACK"
				y.color = "BLACK"
				z.parent.parent.color = "RED"
				z = z.parent.parent
			} else {
				if z == z.parent.right {
					z = z.parent
					leftRotate(T, z)
				}
				z.parent.color = "BLACK"
				z.parent.parent.color = "RED"
				rightRotate(T, z.parent.parent)
			}
		} else { // z.parent == z.parent.parent.right
			y := z.parent.parent.left
			if y != nil && y.color == "RED" {
				z.parent.color = "BLACK"
				y.color = "BLACK"
				z.parent.parent.color = "RED"
				z = z.parent.parent
			} else {
				if z == z.parent.left {
					z = z.parent
					rightRotate(T, z)
				}
				z.parent.color = "BLACK"
				z.parent.parent.color = "RED"
				leftRotate(T, z.parent.parent)
			}
		}
		if z == T.root {
			break
		}
	}
	if T.root != nil {
		T.root.color = "BLACK"
	}
}
func Find(T *RBtree, key uint32) any {
	x := T.root
	for x != T.nil {
		if x.key == key {
			return x.val
		} else if x.key < key {
			x = x.right
		} else {
			x = x.left
		}
	}
	return nil
}

func RBtreeInit() *RBtree {
	return &RBtree{
		root: nil,
		nil:  nil,
	}
}

func NodeInit(key uint32, val any) *Node {
	return &Node{
		key:    key,
		val:    val,
		left:   nil,
		right:  nil,
		parent: nil,
		color:  "RED",
	}
}
